import numpy as np
import tensorflow as tf
from tensorflow.keras import layers, Model
from tensorflow.keras.datasets import cifar10

# 加载CIFAR-10数据集作为示例
(X_train, y_train), (X_test, y_test) = cifar10.load_data()
X_train = X_train.astype('float32') / 255
X_test = X_test.astype('float32') / 255
y_train = y_train.flatten()
y_test = y_test.flatten()

def build_model(input_shape):
    inputs = layers.Input(shape=input_shape)
    x = layers.Conv2D(32, (3, 3), activation='relu')(inputs)
    x = layers.MaxPooling2D((2, 2))(x)
    x = layers.Conv2D(64, (3, 3), activation='relu')(x)
    x = layers.MaxPooling2D((2, 2))(x)
    x = layers.Conv2D(128, (3, 3), activation='relu')(x)
    x = layers.MaxPooling2D((2, 2))(x)
    x = layers.Flatten()(x)
    x = layers.Dense(128, activation='relu')(x)
    outputs = layers.Dense(64, activation=None)(x)  # 输出嵌入向量
    return Model(inputs, outputs)

def triplet_model(input_shape):
    anchor_input = layers.Input(shape=input_shape)
    positive_input = layers.Input(shape=input_shape)
    negative_input = layers.Input(shape=input_shape)
    
    embedding_model = build_model(input_shape)
    
    anchor_embedding = embedding_model(anchor_input)
    positive_embedding = embedding_model(positive_input)
    negative_embedding = embedding_model(negative_input)
    
    model = Model(inputs=[anchor_input, positive_input, negative_input], outputs=[anchor_embedding, positive_embedding, negative_embedding])
    return model

model = triplet_model((32, 32, 3))
model.summary()

def triplet_loss(y_true, y_pred, alpha=0.2):
    anchor_embedding = y_pred[0]
    positive_embedding = y_pred[1]
    negative_embedding = y_pred[2]
    pos_dist = tf.reduce_sum(tf.square(anchor_embedding - positive_embedding), axis=-1)
    neg_dist = tf.reduce_sum(tf.square(anchor_embedding - negative_embedding), axis=-1)
    basic_loss = pos_dist - neg_dist + alpha
    loss = tf.reduce_mean(tf.maximum(basic_loss, 0.0))
    return loss

def create_triplets(X, y, num_triplets):
    triplets = []
    labels_set = set(y)
    label_to_indices = {label: np.where(y == label)[0] for label in labels_set}
    
    for _ in range(num_triplets):
        anchor_label = np.random.choice(list(labels_set))
        negative_label = np.random.choice(list(labels_set - {anchor_label}))
        
        anchor_index = np.random.choice(label_to_indices[anchor_label])
        positive_index = np.random.choice(label_to_indices[anchor_label])
        while positive_index == anchor_index:
            positive_index = np.random.choice(label_to_indices[anchor_label])
        
        negative_index = np.random.choice(label_to_indices[negative_label])
        
        triplets.append([anchor_index, positive_index, negative_index])
    
    return np.array(triplets)

num_triplets = 10000
triplets = create_triplets(X_train, y_train, num_triplets)

def triplet_generator(X, triplets):
    def generator():
        for triplet in triplets:
            anchor_index, positive_index, negative_index = triplet
            anchor = X[anchor_index]
            positive = X[positive_index]
            negative = X[negative_index]
            yield (anchor, positive, negative), np.zeros(3)
    
    return generator

train_generator = triplet_generator(X_train, triplets)

# 使用 tf.data.Dataset.from_generator 创建数据集
train_dataset = tf.data.Dataset.from_generator(
    train_generator,
    output_signature=(
        (tf.TensorSpec(shape=(32, 32, 3), dtype=tf.float32),
         tf.TensorSpec(shape=(32, 32, 3), dtype=tf.float32),
         tf.TensorSpec(shape=(32, 32, 3), dtype=tf.float32)),
        tf.TensorSpec(shape=(3,), dtype=tf.float32)
    )
)

# 设置批量大小
train_dataset = train_dataset.batch(32)

model.compile(optimizer='adam', loss=triplet_loss)
model.fit(train_dataset, epochs=10)

# 保存训练好的模型
model.save('src/model_pth/triplet_model.keras')